#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import json

from fastapi import APIRouter, Query, Request, Form, UploadFile, File
from app.schemas.document_schema import FragmentCutRequest, FragmentDocExtractLocalRequest, FragmentDocExtractUrlRequest
from app.services.document_service.service import DocumentService
from app.services.file_service import FileService
from app.utils.response import success_response, error_response
from app.utils.logging import log_message
from app.services.reranker_service import RerankerService
from typing import List

router = APIRouter(tags=["doc"], prefix="/doc")


@router.post("/cut")
async def txt_cut(
        request: Request,
        payload: FragmentCutRequest  # 直接使用Pydantic模型作为请求体
):
    try:
        service = DocumentService(
            slice_strategy=payload.strategy,
            chunk_size=payload.chunk_size,
            chunk_overlap=payload.chunk_overlap
        )

        # 对文本内容进行切片
        slices = service.process_document(raw_text=payload.texts)

        # 对文档进行重排序
        chunk_number = max(payload.max_length // payload.chunk_size, 1)
        reranker = RerankerService()
        sorted_texts = await reranker.rerank_texts(payload.query, slices["chunks"], chunk_number)

        # 根据输出类型整理数据
        context = "\n\n".join([item["document"]["text"] for item in sorted_texts])
        if payload.chunk:
            chunk = sorted_texts
        else:
            chunk = []

        return success_response({
            "chunk_number": len(sorted_texts),
            "chunk": chunk,
            "context": context
        })
    except Exception as e:
        log_message(
            "doc_error",
            {
                "version": payload.query,  # 从模型获取
                "type": payload.texts,
                "error": str(e)
            },
            "error"
        )
        return error_response(str(e))


@router.post("/smart_extract_local")
async def smart_extract(
        request: Request,
        payload: FragmentDocExtractLocalRequest = Form(...)
):
    """
    第一步，文档上传，可以从用户上传的文件中提取，也可以从用户提供的文件url中提取
    第二步，对文件内容提取
    第三步，对文档进行切片
    第四步，根据用户问题返回，最优切片数据
    """
    try:
        # 第一步：文档上传
        service = DocumentService(
            slice_strategy=payload.strategy,
            chunk_size=payload.chunk_size,
            chunk_overlap=payload.chunk_overlap
        )
        file_service = FileService()
        saved, save_paths = await file_service.save_upload_files(payload.files)

        # 第二步：文档内容提取
        extract_results = service.batch_extract(save_paths)

        # 第三步：对文本进行切片
        slices = [
            f"[file content begin] [{saved[file_path]['original_name']}] {chunk} [file content end]"
            for file_path, file_content in extract_results.items()
            for chunk in service.process_document(file_content)["chunks"]
        ]

        # 第四步：重排序并返回最佳片段
        chunk_number = max(payload.max_length // payload.chunk_size, 1)
        reranker = RerankerService()
        sorted_texts = await reranker.rerank_texts(
            query=payload.query,
            texts=slices,
            top_k=chunk_number
        )

        # 根据输出类型整理数据
        context = "\n\n".join([t["document"]["text"] for t in sorted_texts])

        return success_response({
            "chunk_number": chunk_number,
            "context": context,
            "chunk": sorted_texts if payload.chunk else None,
        })
    except IOError as ioe:
        # 处理文件IO相关错误
        return error_response(json.dumps({
            "type": "file_operation_error",
            "details": str(ioe),
            "files": [f.filename for f in getattr(payload, 'files', [])]
        }))
    except Exception as e:
        log_message(
            "doc_error",
            {
                "version": payload.query,  # 从模型获取
                "error": str(e)
            },
            "error"
        )
        return error_response(str(e))

@router.post("/smart_extract_url")
async def smart_extract(
        request: Request,
        payload: FragmentDocExtractUrlRequest
):
    """
    第一步，文档上传，可以从用户上传的文件中提取，也可以从用户提供的文件url中提取
    第二步，对文件内容提取
    第三步，对文档进行切片
    第四步，根据用户问题返回，最优切片数据
    """
    try:
        # 第一步：文档上传
        service = DocumentService(
            slice_strategy=payload.strategy,
            chunk_size=payload.chunk_size,
            chunk_overlap=payload.chunk_overlap
        )

        file_service = FileService()
        urls = [url.dict() for url in payload.urls]
        saved, save_paths = await file_service.download_and_save_multiple_remote(urls)

        # 第二步：文档内容提取
        extract_results = service.batch_extract(save_paths)

        # 第三步：对文本进行切片
        slices = [
            f"[file content begin] [{saved[file_path]['original_name']}] {chunk} [file content end]"
            for file_path, file_content in extract_results.items()
            for chunk in service.process_document(file_content)["chunks"]
        ]

        # 第四步：重排序并返回最佳片段
        chunk_number = max(payload.max_length // payload.chunk_size, 1)
        reranker = RerankerService()
        sorted_texts = await reranker.rerank_texts(
            query=payload.query,
            texts=slices,
            top_k=chunk_number
        )

        # 根据输出类型整理数据
        context = "\n\n".join([t["document"]["text"] for t in sorted_texts])

        return success_response({
            "chunk_number": chunk_number,
            "context": context,
            "chunk": sorted_texts if payload.chunk else None,
        })
    except IOError as ioe:
        # 处理文件IO相关错误
        return error_response(json.dumps({
            "type": "file_operation_error",
            "details": str(ioe),
            "files": [f.filename for f in getattr(payload, 'files', [])]
        }))
    except Exception as e:
        log_message(
            "doc_error",
            {
                "version": payload.query,  # 从模型获取
                "error": str(e)
            },
            "error"
        )
        return error_response(str(e))
